"""
# @Time    : 2021/7/1 6:53 下午
# @Author  : hezhiqiang01
# @Email   : hezhiqiang01@baidu.com
# @File    : r_actor_critic.py
"""

import torch
import torch.nn as nn
from algorithms.utils.util import init, check
from algorithms.utils.cnn import CNNBase
from algorithms.utils.mlp import MLPBase
from algorithms.utils.rnn import RNNLayer
from algorithms.utils.act import ACTLayer
from algorithms.utils.popart import PopArt
from utils.util import get_shape_from_obs_space


class R_Actor(nn.Module):
    """
    MAPPO的演员网络类。根据观察输出动作。
    :param args: (argparse.Namespace) 包含相关模型信息的参数。
    :param obs_space: (gym.Space) 观察空间。
    :param action_space: (gym.Space) 动作空间。
    :param device: (torch.device) 指定运行设备(cpu/gpu)。
    """
    def __init__(self, args, obs_space, action_space, device=torch.device("cpu")):
        super(R_Actor, self).__init__()
        self.hidden_size = args.hidden_size  # 隐藏层大小

        self._gain = args.gain  # 增益参数
        self._use_orthogonal = args.use_orthogonal  # 是否使用正交初始化
        self._use_policy_active_masks = args.use_policy_active_masks  # 是否在策略中使用活动掩码
        self._use_naive_recurrent_policy = args.use_naive_recurrent_policy  # 是否使用简单循环策略
        self._use_recurrent_policy = args.use_recurrent_policy  # 是否使用循环策略
        self._recurrent_N = args.recurrent_N  # 循环层数
        self.tpdv = dict(dtype=torch.float32, device=device)  # 张量参数字典

        # 根据观察空间形状选择基础网络类型
        obs_shape = get_shape_from_obs_space(obs_space)
        base = CNNBase if len(obs_shape) == 3 else MLPBase
        self.base = base(args, obs_shape)

        # 如果使用循环策略，初始化RNN层
        if self._use_naive_recurrent_policy or self._use_recurrent_policy:
            self.rnn = RNNLayer(self.hidden_size, self.hidden_size, self._recurrent_N, self._use_orthogonal)

        # 初始化动作层
        self.act = ACTLayer(action_space, self.hidden_size, self._use_orthogonal, self._gain)

        # 将模型移动到指定设备
        self.to(device)

    def forward(self, obs, rnn_states, masks, available_actions=None, deterministic=False):
        """
        根据给定输入计算动作。
        :param obs: (np.ndarray / torch.Tensor) 输入网络的观察。
        :param rnn_states: (np.ndarray / torch.Tensor) 如果是RNN网络，则为RNN的隐藏状态。
        :param masks: (np.ndarray / torch.Tensor) 掩码张量，表示是否应将隐藏状态重新初始化为零。
        :param available_actions: (np.ndarray / torch.Tensor) 表示哪些动作对智能体可用
                                                             （如果为None，则所有动作都可用）
        :param deterministic: (bool) 是否从动作分布中采样或返回模式。

        :return actions: (torch.Tensor) 要采取的动作。
        :return action_log_probs: (torch.Tensor) 所采取动作的对数概率。
        :return rnn_states: (torch.Tensor) 更新后的RNN隐藏状态。
        """
        # 将输入转换为张量并移动到指定设备
        obs = check(obs).to(**self.tpdv)
        rnn_states = check(rnn_states).to(**self.tpdv)
        masks = check(masks).to(**self.tpdv)
        if available_actions is not None:
            available_actions = check(available_actions).to(**self.tpdv)

        # 通过基础网络提取特征
        actor_features = self.base(obs)

        # 如果使用循环策略，通过RNN处理特征
        if self._use_naive_recurrent_policy or self._use_recurrent_policy:
            actor_features, rnn_states = self.rnn(actor_features, rnn_states, masks)

        # 计算动作和动作对数概率
        actions, action_log_probs = self.act(actor_features, available_actions, deterministic)

        return actions, action_log_probs, rnn_states

    def evaluate_actions(self, obs, rnn_states, action, masks, available_actions=None, active_masks=None):
        """
        计算给定动作的对数概率和熵。
        :param obs: (torch.Tensor) 输入网络的观察。
        :param action: (torch.Tensor) 要评估熵和对数概率的动作。
        :param rnn_states: (torch.Tensor) 如果是RNN网络，则为RNN的隐藏状态。
        :param masks: (torch.Tensor) 掩码张量，表示是否应将隐藏状态重新初始化为零。
        :param available_actions: (torch.Tensor) 表示哪些动作对智能体可用
                                                （如果为None，则所有动作都可用）
        :param active_masks: (torch.Tensor) 表示智能体是活动的还是死亡的。

        :return action_log_probs: (torch.Tensor) 输入动作的对数概率。
        :return dist_entropy: (torch.Tensor) 给定输入的动作分布熵。
        """
        # 将输入转换为张量并移动到指定设备
        obs = check(obs).to(**self.tpdv)
        rnn_states = check(rnn_states).to(**self.tpdv)
        action = check(action).to(**self.tpdv)
        masks = check(masks).to(**self.tpdv)
        if available_actions is not None:
            available_actions = check(available_actions).to(**self.tpdv)

        if active_masks is not None:
            active_masks = check(active_masks).to(**self.tpdv)

        # 通过基础网络提取特征
        actor_features = self.base(obs)

        # 如果使用循环策略，通过RNN处理特征
        if self._use_naive_recurrent_policy or self._use_recurrent_policy:
            actor_features, rnn_states = self.rnn(actor_features, rnn_states, masks)

        # 评估动作，计算对数概率和熵
        action_log_probs, dist_entropy = self.act.evaluate_actions(actor_features,
                                                                   action, available_actions,
                                                                   active_masks=
                                                                   active_masks if self._use_policy_active_masks
                                                                   else None)

        return action_log_probs, dist_entropy


class R_Critic(nn.Module):
    """
    MAPPO的评论家网络类。根据中心化输入(MAPPO)或局部观察(IPPO)输出价值函数预测。
    :param args: (argparse.Namespace) 包含相关模型信息的参数。
    :param cent_obs_space: (gym.Space) (中心化的)观察空间。
    :param device: (torch.device) 指定运行设备(cpu/gpu)。
    """
    def __init__(self, args, cent_obs_space, device=torch.device("cpu")):
        super(R_Critic, self).__init__()
        self.hidden_size = args.hidden_size  # 隐藏层大小
        self._use_orthogonal = args.use_orthogonal  # 是否使用正交初始化
        self._use_naive_recurrent_policy = args.use_naive_recurrent_policy  # 是否使用简单循环策略
        self._use_recurrent_policy = args.use_recurrent_policy  # 是否使用循环策略
        self._recurrent_N = args.recurrent_N  # 循环层数
        self._use_popart = args.use_popart  # 是否使用PopArt归一化
        self.tpdv = dict(dtype=torch.float32, device=device)  # 张量参数字典
        init_method = [nn.init.xavier_uniform_, nn.init.orthogonal_][self._use_orthogonal]  # 初始化方法

        # 根据观察空间形状选择基础网络类型
        cent_obs_shape = get_shape_from_obs_space(cent_obs_space)
        base = CNNBase if len(cent_obs_shape) == 3 else MLPBase
        self.base = base(args, cent_obs_shape)

        # 如果使用循环策略，初始化RNN层
        if self._use_naive_recurrent_policy or self._use_recurrent_policy:
            self.rnn = RNNLayer(self.hidden_size, self.hidden_size, self._recurrent_N, self._use_orthogonal)

        def init_(m):
            return init(m, init_method, lambda x: nn.init.constant_(x, 0))

        # 根据是否使用PopArt初始化输出层
        if self._use_popart:
            self.v_out = init_(PopArt(self.hidden_size, 1, device=device))
        else:
            self.v_out = init_(nn.Linear(self.hidden_size, 1))

        # 将模型移动到指定设备
        self.to(device)

    def forward(self, cent_obs, rnn_states, masks):
        """
        根据给定输入计算价值。
        :param cent_obs: (np.ndarray / torch.Tensor) 输入网络的观察。
        :param rnn_states: (np.ndarray / torch.Tensor) 如果是RNN网络，则为RNN的隐藏状态。
        :param masks: (np.ndarray / torch.Tensor) 掩码张量，表示是否应将RNN状态重新初始化为零。

        :return values: (torch.Tensor) 价值函数预测。
        :return rnn_states: (torch.Tensor) 更新后的RNN隐藏状态。
        """
        # 将输入转换为张量并移动到指定设备
        cent_obs = check(cent_obs).to(**self.tpdv)
        rnn_states = check(rnn_states).to(**self.tpdv)
        masks = check(masks).to(**self.tpdv)

        # 通过基础网络提取特征
        critic_features = self.base(cent_obs)
        # 如果使用循环策略，通过RNN处理特征
        if self._use_naive_recurrent_policy or self._use_recurrent_policy:
            critic_features, rnn_states = self.rnn(critic_features, rnn_states, masks)
        # 计算价值
        values = self.v_out(critic_features)

        return values, rnn_states
